package com.ruoyi.web.controller.ai;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.aigc.generation.models.QwenParam;
import com.alibaba.dashscope.aigc.imagesynthesis.*;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import io.reactivex.Flowable;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;

import java.util.*;

@RestController
@RequestMapping("/tongyi")
public class TongyiController {

    @Value("${tongyi.apiKey}")
    private String apiKey;

    @GetMapping(value = "/text/{question}", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flowable<GenerationResult> text(@PathVariable String question) throws NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();

        // 创建用户消息对象
        Message userMsg = Message
                .builder()
                .role(Role.USER.getValue())
                .content(question)
                .build();

        // 创建QwenParam对象，设置参数
        QwenParam param = QwenParam.builder()
                .model(Generation.Models.QWEN_MAX)
                .messages(Arrays.asList(userMsg))
                .resultFormat(QwenParam.ResultFormat.MESSAGE)
                .topP(0.8)
                .enableSearch(true)
                .apiKey(apiKey)
                .incrementalOutput(true)    // get streaming output incrementally
                .build();

        // 调用生成接口，获取Flowable对象
        Flowable<GenerationResult> result = gen.streamCall(param);

        return result;

        // 将Flowable转换成Flux<ServerSentEvent<String>>并进行处理
//        return Flux.from(result)
//                // add delay between each event
//                .delayElements(Duration.ofMillis(1000))
//                .map(message -> {
//                    String output = message.getOutput().getChoices().get(0).getMessage().getContent();
//                    System.out.println(output); // print the output
//                    return ServerSentEvent.<String>builder()
//                            .data(output)
//                            .build();
//                })
//                .concatWith(Flux.just(ServerSentEvent.<String>builder().comment("").build()))
//                .doOnError(e -> {
//                    if (e instanceof NoApiKeyException) {
//                        // 处理 NoApiKeyException
//                    } else if (e instanceof InputRequiredException) {
//                        // 处理 InputRequiredException
//                    } else if (e instanceof ApiException) {
//                        // 处理其他 ApiException
//                    } else {
//                        // 处理其他异常
//                        System.out.println("e.getMessage() = " + e.getMessage());
//                    }
//                });
    }


    /**
     * 开始生成任务
     *
     * @param text 文本描述
     * @return
     * @throws NoApiKeyException
     */
    @GetMapping(value = "/image/{text}")
    public List<Map<String, String>> image(@PathVariable String text) throws NoApiKeyException {
        ImageSynthesis is = new ImageSynthesis();
        ImageSynthesisParam param =
                ImageSynthesisParam.builder()
                        .model(ImageSynthesis.Models.WANX_V1)
                        .n(1)
                        .apiKey(apiKey)
                        .size("1024*1024")
                        .prompt(text)
                        .build();
//        ImageSynthesisResult result = is.asyncCall(param);    // 异步返回
        ImageSynthesisResult result = is.call(param);
        /**
         * ImageSynthesisResult(
         * requestId=0579b6c7-99a5-9144-94fb-7844ca9d88c3,
         * output=ImageSynthesisOutput(taskId=d29c70e8-be29-4523-a31d-458e41437a3f,
         * taskStatus=PENDING,
         * code=null, message=null, results=null, taskMetrics=null), usage=null
         * )
         *
         * 任务状态：
         * PENDING 排队中
         * RUNNING 处理中
         * SUCCEEDED 成功
         * FAILED 失败
         * UNKNOWN 作业不存在或状态未知
         */
        List<Map<String, String>> results = result.getOutput().getResults();
        System.err.println(result);
        System.out.println(results);
        return results;
    }

    /**
     * 根据任务id查询任务状态
     * @param taskId    任务id
     * @return
     * @throws NoApiKeyException
     */
    @PostMapping("/image/task/{taskId}")
    public ImageSynthesisOutput task(@PathVariable String taskId) throws NoApiKeyException {
        ImageSynthesis is = new ImageSynthesis();
        ImageSynthesisResult result = is.fetch(taskId, apiKey);
        System.err.println(result);
        return result.getOutput();
    }
}
